import pandas as pd
import redis
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from io import StringIO
import joblib

# 连接Redis获取数据（保持原有数据检查逻辑）
r = redis.Redis(host='localhost', port=6379, decode_responses=True)
table1_json = r.get('japan_trade_table1')
if not table1_json:
    raise ValueError("Redis中没有找到japan_trade_table1数据，请先运行data_processing.py")
    
table1 = pd.read_json(StringIO(table1_json), orient='records')
if table1.empty:
    raise ValueError("从Redis获取的数据为空，请检查数据处理步骤")

table2_json = r.get('japan_trade_table2')
if not table2_json:
    raise ValueError("Redis中没有找到japan_trade_table2数据，请先运行data_processing.py")
    
table2 = pd.read_json(StringIO(table2_json), orient='records')
if table2.empty:
    raise ValueError("从Redis获取的总体数据为空，请检查数据处理步骤")

# 模型1：国家特定预测（输入：年月+国家标签） - 对应API的model1_export/model1_import
valid_dates_mask = table1['日期'].str.match(r'^\d{4}-\d{2}$')
if not valid_dates_mask.any():
    raise ValueError("table1中没有找到格式正确的日期数据（YYYY-MM）")
    
table1 = table1[valid_dates_mask].copy()
table1['年份'] = table1['日期'].str.split('-').str[0].astype(int)
table1['月份'] = table1['日期'].str.split('-').str[1].astype(int)

# 定义特征（必须与API请求的字段完全一致）
X1 = table1[['年份', '月份', '国家标签']]
y1_export = table1['出口额']
y1_import = table1['进口额']

# 预处理管道（与API预测时的输入处理一致）
preprocessor1 = ColumnTransformer(
    transformers=[
        # 国家标签独热编码（处理未知类别）
        ('country_encoder', OneHotEncoder(handle_unknown='ignore'), ['国家标签']),
        # 数值特征标准化（年份/月份）
        ('num_scaler', StandardScaler(), ['年份', '月份'])
    ],
    verbose_feature_names_out=False  # 避免特征名带前缀，保持清晰
)

# 训练国家特定出口模型（保存为model1_export.pkl）
model1_export = Pipeline([
    ('preprocessor', preprocessor1),
    ('regressor', LinearRegression())
])
model1_export.fit(X1, y1_export)

# 训练国家特定进口模型（保存为model1_import.pkl）
model1_import = Pipeline([
    ('preprocessor', preprocessor1),
    ('regressor', LinearRegression())
])
model1_import.fit(X1, y1_import)

# 模型2：总体预测（输入：年月） - 对应API的model2_export/model2_import
valid_dates_mask = table2['日期'].str.match(r'^\d{4}-\d{2}$')
if not valid_dates_mask.any():
    raise ValueError("table2中没有找到格式正确的日期数据（YYYY-MM）")
    
table2 = table2[valid_dates_mask].copy()
table2['年份'] = table2['日期'].str.split('-').str[0].astype(int)
table2['月份'] = table2['日期'].str.split('-').str[1].astype(int)

X2 = table2[['年份', '月份']]  # 特征与API请求字段一致
y2_export = table2['出口额']
y2_import = table2['进口额']

# 总体模型预处理（仅标准化数值特征）
preprocessor2 = ColumnTransformer(
    transformers=[
        ('num_scaler', StandardScaler(), ['年份', '月份'])
    ],
    verbose_feature_names_out=False
)

# 训练总体出口模型（保存为model2_export.pkl）
model2_export = Pipeline([
    ('preprocessor', preprocessor2),
    ('regressor', LinearRegression())
])
model2_export.fit(X2, y2_export)

# 训练总体进口模型（保存为model2_import.pkl）
model2_import = Pipeline([
    ('preprocessor', preprocessor2),
    ('regressor', LinearRegression())
])
model2_import.fit(X2, y2_import)

# 保存模型（与API加载路径完全一致）
joblib.dump(model1_export, 'model1_export.pkl')
joblib.dump(model1_import, 'model1_import.pkl')
joblib.dump(model2_export, 'model2_export.pkl')
joblib.dump(model2_import, 'model2_import.pkl')

print("模型训练完成，已保存为model1_export.pkl、model1_import.pkl、model2_export.pkl、model2_import.pkl")
